# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
import pprint

from absl import app, flags

from coin.base.datetime_util import to_datetime
from coin.base.timestamp import get_timestamp

from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct

from coin.exchange.base.order_gateway import OrderSide
from coin.strategy.mm.base.dump.pass_unhedge import FillRow, FeatureRow
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.mm.base.pass_unhedge import Pricer
from coin.strategy.mm.base.status_printer import PeriodicStatusPrinter

from coin.strategy.mm.bitmex.simple_aggressor import BitmexSimpleAggressor
from coin.strategy.mm.bitmex.simple_mm_executor import BitmexSimpleMMExecutor
from coin.strategy.mm.bitmex.stop_limit_mm_executor import BitmexStopLimitMMExecutor
from coin.strategy.mm.bitmex.stop_loss_executor import BitmexStopLossExecutor

from coin.strategy.mm.bitmex.pass_unhedge_param import get_strategy_params
from coin.strategy.mm.base.bitmex_strategy_base import (BitmexPositionBalancePrinter,
                                                        BitmexStrategyBase,
                                                        run_live)

from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.proto.coin_request_pb2 import StrategyRequestProto


class PassiveUnhedgeStrategy(BitmexStrategyBase):
  def __init__(self, strategy_params, dry_run=True):
    trade_product = strategy_params['trade_product']
    super().__init__(products=[trade_product], logger=logging.getLogger('PassiveUnhedgeStrategy'))

    self._strategy_params = strategy_params
    self._dry_run = dry_run
    self._shutdown = False

    self._trade_product = trade_product
    self._ref_product = strategy_params['ref_product']

    self._edge = strategy_params['edge_bp'] / 10000.
    self._close_edge = strategy_params['close_edge_bp'] / 10000.
    if strategy_params.get('agg_edge_bp', None) is None:
      self._agg_edge = None
    else:
      self._agg_edge = strategy_params['agg_edge_bp'] / 10000.

    self._min_pos = float(strategy_params['min_pos'])
    self._max_pos = float(strategy_params['max_pos'])

    self._order_update_period = (strategy_params.get('order_update_period', 2) * (10**9))
    self._enable_queue0_executor = (strategy_params.get('enable_queue0_executor', False))
    self._enable_stop_loss_executor = (strategy_params.get('enable_stop_loss_executor', False))

    self._bookmap = BookMap([self._trade_product, self._ref_product], logger=self._logger)
    self._pricer = Pricer(**strategy_params['pricing'])
    self._pass_executor = None
    self._queue0_executor = None
    self._stop_loss_executor = None
    self._agg_executor = None

    self._trade_begin_ts = None
    self._last_fire_ts = 0

    request = StrategyRequestProto(**self._strategy_params['logger'])
    self._strategy_logger = StrategyLogProvider(request)
    self._strategy_logger.write_request()

    self._printer = BitmexPositionBalancePrinter(self._trade_product,
                                                 self._bookmap,
                                                 strategy_logger=self._strategy_logger,
                                                 logger=self._logger)
    self._status_printer = PeriodicStatusPrinter(logger=self._logger)
    self._status_printer.add_status_update_callback(self._update_status)
    self._last_feature_dump_ts = 0

    # Print header
    self._logger.info(FillRow.get_header())
    self._logger.info(FeatureRow.get_header())

  @property
  def strategy_logger(self):
    return self._strategy_logger

  def prepare(self):
    assert self.feed_subsystem is not None

    # Subscribe feed.
    for product in (self._trade_product, self._ref_product):
      self.feed_subsystem.subscribe_book(product, functools.partial(self.on_book, product))

  def on_og_reset(self):
    # Start balance and position printing.
    self._printer.register_order_gateway(self.order_gateway)
    self._printer.print_position_balance(ignore_error=True)
    self._printer.start_print(self._ioloop, 10)  # 10 seconds

    self._status_printer.print_status(ignore_error=True)
    self._status_printer.start(10000)  # 10 seconds

    # Intialize passive executor
    self._logger.info('Intializing executor...')

    self._pass_executor = BitmexSimpleMMExecutor(self._strategy_params['pass_executor'],
                                                 self.order_gateway)

    if self._enable_queue0_executor:
      self._queue0_executor = BitmexStopLimitMMExecutor(self._strategy_params['queue0_executor'],
                                                        self.order_gateway)

    if self._enable_stop_loss_executor:
      self._stop_loss_executor = BitmexStopLossExecutor(self._strategy_params['stop_loss_executor'],
                                                        self.order_gateway)

    self._agg_executor = BitmexSimpleAggressor(self._strategy_params['agg_executor'],
                                               self.order_gateway)
    self.order_gateway.register_fill_callback(self._on_fill)

  def on_book(self, product, book):
    self._bookmap.set_book(product, book)

    # Check book freshness and sanity
    if not self._bookmap.is_book_ready() or not self._bookmap.is_book_sane():
      return

    self._pricer.update(self._bookmap.get_book(self._ref_product),
                        self._bookmap.get_book(self._trade_product))

    if (not self._pricer.ready or not self.is_og_ready() or self._dry_run or self._shutdown):
      return

    now_ts = get_timestamp()
    if self._trade_begin_ts is None:
      self._trade_begin_ts = now_ts

    if now_ts - self._last_fire_ts >= self._order_update_period:
      create_list = []
      amend_list = []
      self.execute_agg(create_list, amend_list)
      self.execute_pass(create_list, amend_list)
      self.execute_queue0(create_list, amend_list)
      self.execute_stop_loss(create_list, amend_list)

      self.execute(create_list, amend_list)
      self._last_fire_ts = now_ts

    if book.timestamp - self._last_feature_dump_ts >= 2 * (10**9):
      self._last_feature_dump_ts = book.timestamp
      self._logger.info(self._pricer.last_value.to_log_string())

  def execute(self, create_list, amend_list):
    if amend_list:
      self.order_gateway.amend_bulk(amend_list)
      self._logger.debug('Amend %d orders', len(amend_list))
    if create_list:
      self.order_gateway.submit_bulk(create_list)
      self._logger.debug('Create %d orders', len(create_list))

  def get_working_order_price(self):
    assert self.is_og_ready()
    askp = None
    bidp = None

    for order in self.order_gateway.get_working_order():
      if order.order_side == OrderSide.SELL:
        if askp is None or askp > order.price:
          askp = order.price
      elif order.order_side == OrderSide.BUY:
        if bidp is None or bidp < order.price:
          bidp = order.price
    return askp, bidp

  def execute_agg(self, create_list, amend_list):
    if self._agg_edge is None:
      return
    assert self._agg_executor is not None

    self._agg_executor.cancel_all(amend_list=amend_list, fire=False)

    askp, bidp = self._pricer.calculate_price(self._agg_edge, self._agg_edge, False)

    book = self._bookmap.get_book(self._trade_product)
    best_askp, best_bidp = self.get_working_order_price()

    # Do not buy and sell at the same time.
    if askp <= book.bid0().price:
      if best_bidp is not None:
        # To prevent self-trading, keep the agg price above exisiting buy order.
        askp = max(self._pricer.tick.get_next_price(best_bidp), askp)
      self._agg_executor.sell(askp, create_list=create_list, fire=False)
    elif bidp >= book.ask0().price:
      if best_askp is not None:
        # To prevent self-trading, keep the agg price below exisiting sell order.
        bidp = min(self._pricer.tick.get_prev_price(best_askp), bidp)
      self._agg_executor.buy(bidp, create_list=create_list, fire=False)

  def execute_pass(self, create_list, amend_list):
    pos = self.order_gateway.get_position(self._trade_product)
    sell_edge = self._edge
    if pos >= self._max_pos:
      sell_edge = self._close_edge
    elif pos >= 0:
      p = pos / self._max_pos
      assert 0. <= p <= 1.
      sell_edge = (1. - p) * self._edge + p * self._close_edge

    buy_edge = self._edge
    if pos <= self._min_pos:
      buy_edge = self._close_edge
    elif pos <= 0:
      p = pos / self._min_pos
      assert 0. <= p <= 1.
      buy_edge = (1. - p) * self._edge + p * self._close_edge

    askp, bidp = self._pricer.calculate_price(sell_edge, buy_edge, True)
    self._pass_executor.manage_mm_orders(askp, bidp, create_list, amend_list, fire=False)

  def execute_queue0(self, create_list, amend_list):
    if not self._enable_queue0_executor:
      return
    assert self._queue0_executor is not None

    book = self._bookmap.get_book(self._trade_product)

    pos = self.order_gateway.get_position(self._trade_product)
    sell_edge = self._edge
    if pos >= self._max_pos:
      sell_edge = self._close_edge
    elif pos >= 0:
      p = pos / self._max_pos
      assert 0. <= p <= 1.
      sell_edge = (1. - p) * self._edge + p * self._close_edge

    buy_edge = self._edge
    if pos <= self._min_pos:
      buy_edge = self._close_edge
    elif pos <= 0:
      p = pos / self._min_pos
      assert 0. <= p <= 1.
      buy_edge = (1. - p) * self._edge + p * self._close_edge

    askp, bidp = self._pricer.calculate_price(sell_edge, buy_edge, False)
    self._queue0_executor.manage_orders(book.ask0().price,
                                        book.bid0().price,
                                        askp,
                                        bidp,
                                        create_list=create_list,
                                        amend_list=amend_list,
                                        fire=False)

  def execute_stop_loss(self, create_list, amend_list):
    if not self._enable_stop_loss_executor:
      return
    assert self._stop_loss_executor is not None

    book = self._bookmap.get_book(self._trade_product)

    self._stop_loss_executor.manage_orders(book.ask0().price,
                                           book.bid0().price,
                                           create_list=create_list,
                                           amend_list=amend_list,
                                           fire=False)

  def shutdown(self):
    if self._shutdown:
      return
    self._shutdown = True

    if not self.is_og_ready():
      self._logger.info('Shutting down now')
      self._ioloop.stop()
      return

    self._ioloop.add_timeout(datetime.timedelta(seconds=3), self._ioloop.stop)
    self._logger.info('Shutting down! Canceling all working orders and '
                      'no new order will be submitted. Exit in 3 seconds.')

    # Cancel all now and again 1 second later.
    self.order_gateway.cancel_all()
    self._ioloop.add_timeout(datetime.timedelta(seconds=1), self.order_gateway.cancel_all)
    self._ioloop.add_timeout(datetime.timedelta(seconds=2), self._status_printer.print_status)

  def _on_fill(self, fill_price, fill_qty, order):
    timestamp = get_timestamp()

    pos, unrealized_pnl = self._printer._get_pos(self._trade_product)
    balance = self.order_gateway.get_balance(self._trade_product.quote)
    balance += unrealized_pnl

    if order.order_side == OrderSide.BUY:
      side_sign = 1
    elif order.order_side == OrderSide.SELL:
      side_sign = -1
    else:
      raise ValueError('Invalid order side: %s' % order.order_side)

    pos += side_sign * fill_qty
    fee = order.internal.commission or -1
    row = FillRow(time=timestamp,
                  time_hr=to_datetime(timestamp),
                  exchange='Bitmex',
                  product=str(self._trade_product),
                  side=side_sign,
                  price=fill_price,
                  qty=fill_qty,
                  maker=-1,
                  fee=fee,
                  pos=pos,
                  balance=balance)
    self._logger.info(row.to_log_string())

  def _update_status(self, update_time, status_out):
    # Static status
    if 'name' not in status_out:
      status_out['name'] = self._strategy_params.get('name', None)
      status_out['log_name'] = self._strategy_params.get('logger', {}).get('strategy_name', None)
      status_out['group'] = self._strategy_params.get('group', None)
      status_out['account'] = self._strategy_params.get('account', None)
      status_out['pass_unhedge.trade_product'] = str(self._trade_product)
      status_out['pass_unhedge.ref_product'] = str(self._ref_product)
      status_out['pass_unhedge.basis_ma_window_min'] = '%.1f' % (
          self._strategy_params['pricing']['basis_ma_window'] / (60 * 10**9))
      status_out['pass_unhedge.edge_bp'] = '%.1f' % (self._edge * 10000)
      status_out['pass_unhedge.close_edge_bp'] = '%.1f' % (self._close_edge * 10000)
      status_out['pass_unhedge.agg_edge_bp'] = (None if self._agg_edge is None else
                                                ('%.1f' % (self._agg_edge * 10000)))

      lot_size = self._strategy_params['pass_executor']['lot_size']
      max_pos = self._strategy_params['pass_executor']['max_pos']
      status_out['pass_unhedge.lot_size'] = lot_size
      status_out['pass_unhedge.stack'] = 2 * max_pos // lot_size

    if ('trade_begin_time' not in status_out and self._trade_begin_ts is not None):
      status_out['trade_begin_time'] = str(to_datetime(self._trade_begin_ts))

    # Pricer
    if self._pricer is not None and self._pricer.ready:
      status_out['pricer.ask_basis_ma'] = '%.4f' % self._pricer.last_value.ask_basis_ma
      status_out['pricer.bid_basis_ma'] = '%.4f' % self._pricer.last_value.bid_basis_ma
      window_len = self._pricer.window_len
      status_out['pricer.ask_basis_window_len'] = window_len[0]
      status_out['pricer.bid_basis_window_len'] = window_len[1]

    # Position and balance
    self._printer.update_position_balance_status(status_out)


def get_ref_product(exchange):
  if exchange == 'Okex':
    ref_product = OkexFuturesProduct.FromStr('BTC-USD.QUARTER')
    ref_product_true_book_funcs = ((lambda book: book.get_notional_asks_by_qty(100.)[1][-1][0]),
                                   (lambda book: book.get_notional_bids_by_qty(100.)[1][-1][0]))

  elif exchange == 'Bitflyer':
    ref_product = BitflyerFuturesProduct.FromStr('BTC-JPY.PERPETUAL')
    ref_product_true_book_funcs = ((lambda book: book.get_notional_asks_by_qty(1.)[1][-1][0]),
                                   (lambda book: book.get_notional_bids_by_qty(1.)[1][-1][0]))

  elif exchange == 'Binance':
    ref_product = BinanceProduct.FromStr('BTC-USDT')
    ref_product_true_book_funcs = ((lambda book: book.get_notional_asks_by_qty(1.)[1][-1][0]),
                                   (lambda book: book.get_notional_bids_by_qty(1.)[1][-1][0]))

  else:
    raise ValueError(exchange)

  return ref_product, ref_product_true_book_funcs


def main(argv):
  FLAGS = flags.FLAGS

  assert len(argv) in [2, 3]
  target = argv[1]
  og_config = None if len(argv) < 3 else argv[2]

  logging.info('target=%s, og_config=%s', target, og_config)

  strategy_params, env_config = get_strategy_params(target,
                                                    group=FLAGS.group,
                                                    tag=FLAGS.tag,
                                                    og_config=og_config,
                                                    init_og_immediately=FLAGS.no_warming_up)

  logging.info('strategy_params\n%s', pprint.pformat(strategy_params, width=120))
  logging.info('env_config\n%s', pprint.pformat(env_config, width=120))

  # Run strategy!
  strategy = PassiveUnhedgeStrategy(strategy_params, dry_run=False)
  return run_live(strategy, **env_config)


if __name__ == '__main__':
  logging.getLogger('requests').setLevel(logging.WARNING)
  logging.getLogger('urllib3').setLevel(logging.WARNING)
  logging.getLogger('coin.exchange.base.rest.private_queryer_base').setLevel(logging.WARNING)
  logging.getLogger('coin.exchange.okex_futures.ws.subscriber').setLevel(logging.WARNING)

  flags.DEFINE_string('group', None, 'group name')
  flags.DEFINE_string('tag', None, 'If specified, it will be appended to strategy name')
  flags.DEFINE_bool('no_warming_up',
                    False,
                    'If true, it starts to trade immediately without warming up.')
  app.run(main)
